!test

!test

Greptile Summary

This PR implements the auto TMA transpose scheduler, introducing a new TmaTranspose enable option and replacing the previous stub with a fully functional getTransposeHeuristics / scheduleTranspose pair in transpose_tma.cpp. The scheduler selects between two transpose strategies based on input/output tensor counts: when there are more inputs than outputs (is_output_smem_transpose = true), TMA store with swizzle is used on the output side; otherwise TMA load with swizzle operates on the input side. Supporting changes include new fields on TransposeParams, a C++20 ranges refactor in tma.cpp, and std::uint8_t underlying types on all option enums.

Key issues found:

• Infinite loop in getTransposeHeuristics: if bytes_per_cta / n_input < kTmaSwizzleBytes (128), integer division produces estimated_tile_size1 = 0; the while loop multiplying by 2 never converges (0 * 2 == 0). This can be triggered when n_input > 64 on certain GPU configs.
• Missing guard in scheduleTranspose: tma_store_tvs.at(0) is accessed whenever is_output_smem_transpose == true without first asserting use_tma_store == true, unlike the symmetric check already present for the input-smem path.
• Debug std::cout left in OutputTransposeBankconflict test — fires on test failure and should use GTest facilities instead.
• Typos ("tranapose" → "transpose") in comment lines describing test parameter combinations.

Confidence Score: 3/5

• PR is functionally correct for typical 1–2 input transpose fusions, but contains an infinite-loop risk and a potential crash for less common parameter combinations.
• The PR introduces the TMA transpose scheduler with correct overall logic for standard use cases (1–2 inputs). However, two logic issues lower confidence:

1. Infinite loop when estimated_tile_size1 = 0: This occurs when n_input > 64 on SM90, a realistic edge case for more complex fusion patterns. The while loop condition guarantees termination only if estimated_tile_size1 > 0.

2. Unguarded tma_store_tvs.at(0) access: While the default heuristic ties is_output_smem_transpose and use_tma_store together, the API allows independent override (as shown in parameterized tests). Missing the guard that exists in the symmetric input-smem path is an asymmetry that could cause crashes if those params are misused.

Neither issue is triggered by the test suite as written, but both represent real failure modes for edge-case inputs or manually-constructed params. The code quality issues (test cout and typos) are minor.

• csrc/scheduler/transpose_tma.cpp — review the tile-size heuristic initialization and the is_output_smem_transpose / use_tma_store constraint in scheduleTranspose.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[TransposeScheduler::computeHeuristics] --> B{TmaTranspose enabled?}
    B -- Yes --> C[tma::getTransposeHeuristics]
    B -- No --> E[non_tma::getTransposeHeuristics]
    C --> D{n_input > n_output?}
    D -- Yes\nis_output_smem_transpose=true --> F[use_tma_load=true\nuse_tma_store=true\nswizzle on output smem]
    D -- No\nis_output_smem_transpose=false --> G[use_tma_load=true\nuse_tma_store=false\nswizzle on input smem]
    F --> H[Return TransposeParams]
    G --> H
    C -- returns null --> E
    E --> H
    H --> I[TransposeScheduler::schedule]
    I --> J{use_tma_load OR use_tma_store?}
    J -- Yes --> K[tma::scheduleTranspose]
    J -- No --> L[non_tma::scheduleTranspose]
    K --> M{is_output_smem_transpose?}
    M -- true --> N[TMA load w/o swizzle\nTMA store w/ MmaSwizzle\ntranspose at regs→output smem]
    M -- false --> O[TMA load w/ MmaSwizzle\nregister store\ntranspose at input smem→regs]

_{Last reviewed commit: bc772db}

Review updated until commit bc772db

Description

• Implements automatic TMA (Tensor Memory Access) transpose scheduler with two paths: input smem transpose (swizzle on input) and output smem transpose (swizzle on output)

• Adds new TmaTranspose enable option to toggle the feature; scheduler falls back to non-TMA when disabled

• Introduces new parameters: use_tma_store, is_output_smem_transpose, chunks_per_thread, elements_per_chunk for flexible TMA configuration

• Adds comprehensive tests covering different dtypes, transpose dimensions, and TMA parameter combinations

auto non_trivial_ids =
    tv->getLoopDomain() | std::views::filter([](const IterDomain* id) {
      return !id->extent()->isConstScalar() ||
          id->extent()->evaluate().as<int64_t>() > 1;
    });
if (std::ranges::any_of(non_trivial_ids, [](const IterDomain* id) {
      return id->isThreadDim() ||
          id->getParallelType() == ParallelType::Serial;
    })) {
  return {};
}

for (auto [cached_output, output_idx] : cached_outputs) {
  auto output = fusion->outputs()[output_idx]->as<TensorView>();
  output->definition()->as<LoadStoreOp>()->setOpType(
      LoadStoreOpType::CpAsyncBulkTensorTile);
  cached_output->setMemoryType(MemoryType::Shared);
  cached_output->cacheBefore();
  tma_store_tvs.push_back(cached_output);
}

Options(const Options& other)
    : options_([&other]() {
        std::lock_guard<std::mutex> lock_other(other.mutex_);
        return other.options_;
      }()) {}